import numpy as np
import os


def gen_l1loss_grad_2d_data():
    """
    为 L1LossGrad 算子生成完整测试数据（对齐 PyTorch 行为）：
      1. 输入  : input_x.bin、input_y.bin
      2. 上游梯度 : grad_output_none.bin / grad_output_sum.bin / grad_output_mean.bin
                   （已广播成 [B,F] tensor，kernel 端直接 Mul 即可）
      3. Golden : golden_*.bin
      4. 模式标记 : reduction_mode.txt（none,sum,mean）
    固定种子保证可重复。
    """
    np.random.seed(42)
    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)

    # 1. 形状常量
    batch_size = 13
    features   = 1525
    numel      = batch_size * features

    # 2. 随机生成预测值 x 与目标值 y
    input_pred   = np.random.uniform(-1, 1, (batch_size, features)).astype(np.float32)
    input_target = np.random.uniform(-1, 1, (batch_size, features)).astype(np.float32)

    # 3. 计算 sign(delta)
    sign_delta = np.sign(input_pred - input_target).astype(np.float32)

    # 4. 上游梯度 ---- 按 PyTorch 约定 -----------------
    grad_output_none = np.random.uniform(-1, 1, (batch_size, features)).astype(np.float32)

    grad_output_sum_tensor  = np.full((batch_size, features), 1.0, dtype=np.float32)          # sum 模式
    grad_output_mean_tensor = np.full((batch_size, features), 1.0 / numel, dtype=np.float32)  # mean 模式

    # 5. 计算 golden
    golden_none = grad_output_none * sign_delta
    golden_sum  = grad_output_sum_tensor * sign_delta
    golden_mean = grad_output_mean_tensor * sign_delta

    # 6. 落盘输入
    input_pred.tofile("./input/input_x.bin")
    input_target.tofile("./input/input_y.bin")

    
    # 7. 落盘上游梯度
    grad_output_none.tofile("./input/grad_output_none.bin")
    grad_output_sum_tensor.tofile("./input/grad_output_sum.bin")
    grad_output_mean_tensor.tofile("./input/grad_output_mean.bin")

    # 8. 落盘 golden
    golden_none.tofile("./output/golden_none.bin")
    golden_sum.tofile("./output/golden_sum.bin")
    golden_mean.tofile("./output/golden_mean.bin")

    # 10. 抽查
    print("=== 前16个值 ===")
    print("sign_delta :", sign_delta.flatten()[:16])
    print("golden_none:", golden_none.flatten()[:16])
    print("golden_sum :", golden_sum.flatten()[:16])
    print("golden_mean:", golden_mean.flatten()[:16])
    print("\n=== 统计信息 ===")
    print("input_x shape:", input_pred.shape)
    print("input_y shape:", input_target.shape)
    print("numel:", numel)
    print("mean 上游梯度系数:", 1.0 / numel)


if __name__ == "__main__":
    gen_l1loss_grad_2d_data()